package services

import (
	"context"
	"sync"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/service/ec2"
)

type AWSClientManager struct {
	clients map[string]*ec2.Client
	mu      sync.RWMutex
}

func NewAWSClientManager() *AWSClientManager {
	return &AWSClientManager{
		clients: make(map[string]*ec2.Client),
	}
}

func (m *AWSClientManager) GetEC2Client(region string) (*ec2.Client, error) {
	m.mu.RLock()
	if client, exists := m.clients[region]; exists {
		m.mu.RUnlock()
		return client, nil
	}
	m.mu.RUnlock()

	// 创建新的客户端
	cfg, err := config.LoadDefaultConfig(context.TODO(),
		config.WithRegion(region),
	)
	if err != nil {
		return nil, err
	}

	client := ec2.NewFromConfig(cfg)

	m.mu.Lock()
	m.clients[region] = client
	m.mu.Unlock()

	return client, nil
}

func (m *AWSClientManager) GetAllRegions() ([]string, error) {
	// 使用默认区域获取所有可用区域
	client, err := m.GetEC2Client("us-east-1")
	if err != nil {
		return nil, err
	}

	result, err := client.DescribeRegions(context.TODO(), &ec2.DescribeRegionsInput{
		AllRegions: aws.Bool(false),
	})
	if err != nil {
		return nil, err
	}

	regions := make([]string, len(result.Regions))
	for i, region := range result.Regions {
		regions[i] = *region.RegionName
	}

	return regions, nil
}
